from mindspore import load_checkpoint, load_param_into_net
from mindspore import Model
from mindspore.nn import Accuracy
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
from mindspore.train.callback import Callback

import os

class StepLossAccInfo(Callback):
    def __init__(self, model, eval_dataset, steps_loss, steps_eval):
        self.model = model
        self.eval_dataset = eval_dataset
        self.steps_loss = steps_loss
        self.steps_eval = steps_eval

    def step_end(self, run_context):
        cb_params = run_context.original_args()
        cur_epoch = cb_params.cur_epoch_num
        cur_step = (cur_epoch - 1) * 1875 + cb_params.cur_step_num
        self.steps_loss["loss_value"].append(str(cb_params.net_outputs))
        self.steps_loss["step"].append(str(cur_step))
        if cur_step % 125 == 0:
            acc = self.model.eval(self.eval_dataset, dataset_sink_mode=False)
            self.steps_eval["step"].append(cur_step)
            self.steps_eval["acc"].append(acc["Accuracy"])


def load_model(network, model_file):
    param_dict = load_checkpoint(model_file)
    # load parameter to the network
    load_param_into_net(network, param_dict)


def train(network,net_opt, net_loss, model_path, epoch_num, ds_train, ds_eval=None):

    # clean up old run files before in Linux
    os.system('rm -f {0}*.ckpt {0}*.meta {0}*.pb'.format(model_path))

    # define the model
    model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()})
    config_ck = CheckpointConfig(save_checkpoint_steps=375, keep_checkpoint_max=5)
    # group layers into an object with training and evaluation features
    ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet", directory=model_path, config=config_ck)

    steps_loss = {"step": [], "loss_value": []}
    steps_eval = {"step": [], "acc": []}
    # collect the steps,loss and accuracy information
    step_loss_acc_info = StepLossAccInfo(model, ds_eval, steps_loss, steps_eval)

    print("============== Starting Training ==============")
    model.train(epoch_num, ds_train, callbacks=[ckpoint_cb, LossMonitor(125), step_loss_acc_info],
                dataset_sink_mode=False)


def eval(network, model_file, net_loss, ds_eval):

    param_dict = load_checkpoint(model_file)
    # load parameter to the network
    load_param_into_net(network, param_dict)
    model = Model(network, net_loss, metrics={"Accuracy": Accuracy()})

    print("============== Starting Testing ==============")
    acc = model.eval(ds_eval, dataset_sink_mode=False)
    print("============== Accuracy:{} ==============".format(acc))
    pass